use std::io::Write;
use std::path::Path;
use std::fs::{self, File};



// ------------ Parser -------------------

#[derive(Debug)]
enum Segment {
    Argument, Local, That, This,
    Pointer, Temp,
    Constant,
    Static,
}

#[derive(Debug)]
enum Code {
    Add, Sub, Neg, 
    Eq,  Gt,  Lt,
    And, Or,  Not,
    Push(Segment, u16),
    Pop(Segment, u16),
}

use Code::*;
use Segment::*;
struct Parser { content: String }

impl Parser {

    fn new(file: &str) -> Self {
        let content = fs::read_to_string(file).unwrap();
        Parser { content }
    }

    fn parse(self) -> Vec<Code> {
        let code = self.filter_comment_line();
        code.into_iter()
            .map(|line| Self::parse_single_line(line))
            .collect()
    }

    fn filter_comment_line(self) -> Vec<String> {
        let codes = self.content.split("\n")
            .filter(|x| !(x.starts_with("//") || x.trim() == ""))
            .map(|x| x.to_string())
            .collect();
        return codes;
    }

    fn parse_single_line(line: String) -> Code {
        let strlist: Vec<&str> = line.as_str().trim().split(" ").collect();
        match strlist[0] {
            "add" => Add, "sub" => Sub, "neg" => Neg,
            "eq"  => Eq,  "gt"  => Gt,  "lt"  => Lt,
            "and" => And, "or"  => Or,  "not" => Not,
            "push" | "pop" => {
                let segm = match strlist[1] {
                    "local" => Local,
                    "argument" => Argument,
                    "that" => That,
                    "this" => This,
                    "pointer" => Pointer,
                    "temp" => Temp,
                    "constant" => Constant,
                    "static" => Static,
                    _ => panic!("Invalid Segment!"),
                };
                let index: u16 = strlist[2].parse().unwrap();
                if strlist[0] == "push" {
                    return Push(segm, index);
                }
                return Pop(segm, index);
            }
            _ => panic!("Invalid code"),
        }
    }
}

// ------------ Code Writer -------------------

fn handle_comp(comp: &str, idx: usize) -> String {
    let res = "@SP    \n".to_owned()
            + "AM=M-1 \n"
            + "D=M    \n"
            + "A=A-1  \n"
            + "D=M-D  \n";
    // jump here
    let true_i = format!("TRUE{}", idx);
    let continue_i = format!("CONTINUE{}", idx);
    let jump = format!( 
       "@{}
        D;{}
        D=0
        @{}
        0;JMP
    ({})
        D=-1
    ({})
        @SP
        AM=M-1
        M=D
        @SP
        M=M+1
    ", true_i, comp, continue_i, true_i, continue_i);
    return res + &jump;
}



struct CodeWriter {}
impl CodeWriter {
    fn generate_single_code(fname: &str, code: Code, idx: usize) -> String {
        match code {
            Push(segm, index) => {
                let offset = format!("@{}   \n", index)
                                   + "D=A   \n";
                let set_data = match segm {
                    Constant => offset,
                    Argument => offset + "@ARG   \n  A=D+M  \n D=M \n",
                    Local    => offset + "@LCL   \n  A=D+M  \n D=M \n",
                    This     => offset + "@THIS  \n  A=D+M  \n D=M \n",
                    That     => offset + "@THAT  \n  A=D+M  \n D=M \n",
                    Pointer  => offset + "@3     \n  A=D+A  \n D=M \n",
                    Temp     => offset + "@5     \n  A=D+A  \n D=M \n",
                    Static   => format!( "@{}.{} \n", fname, index) + "D=M \n",
                };
                let in_stack = "@SP   \n\
                                A=M   \n\
                                M=D   \n\
                                @SP   \n\
                                M=M+1 \n";
                return set_data + in_stack;
            }
            Pop(segm, index) => {
                let offset = format!("@{}   \n", index)
                                   + "D=A   \n";
                let set_data = "@R13    \n\
                                M=D     \n\
                                @SP     \n\
                                AM=M-1  \n\
                                D=M     \n\
                                @R13    \n\
                                A=M     \n\
                                M=D     \n";
                match segm {
                    Argument => offset + "@ARG   \n  D=D+M   \n" + set_data,
                    Local    => offset + "@LCL   \n  D=D+M   \n" + set_data,
                    This     => offset + "@THIS  \n  D=D+M   \n" + set_data,
                    That     => offset + "@THAT  \n  D=D+M   \n" + set_data,
                    Pointer  => offset + "@3     \n  D=D+A   \n" + set_data,
                    Temp     => offset + "@5     \n  D=D+A   \n" + set_data,
                    Static   => format!( "@SP    \n\
                                          AM=M-1 \n\
                                          D=M    \n\
                                          @{}.{} \n\
                                          M=D    \n", fname, index),
                    Constant => panic!("Could not pop on Constant Segment"),
                }
            }
            Add | Sub | Or | And => {
                let c1 = "@SP     \n".to_owned() 
                       + "AM=M-1  \n"
                       + "D=M     \n"
                       + "A=A-1   \n";
                let c2 = match code {
                    Add => "M=D+M \n",
                    Sub => "M=M-D \n",
                    Or  => "M=D|M \n",
                    And => "M=D&M \n",
                    _ => unreachable!(),
                };
                return c1 + c2;
            }
            Neg | Not => {
                let c1 = "@SP   \n".to_owned() 
                       + "AM=M-1\n";
                let c2 = match code {
                    Neg => "M=-M\n",
                    Not => "M=!M\n",
                    _ => unreachable!(),
                };
                let c3 = "@SP   \n";
                let c4 = "M=M+1 \n";
                return c1 + c2 + c3 + c4;
            }
            Eq => handle_comp("JEQ", idx),
            Gt => handle_comp("JGT", idx),
            Lt => handle_comp("JLT", idx),
        }
    }

    fn generate_code(fname: &str, output: &str, codes: Vec<Code>) -> std::io::Result<()> {
        let mut file = File::create(output)?;
        for (i, code) in codes.into_iter().enumerate() {
            let code_ = Self::generate_single_code(fname, code, i);
            file.write(code_.as_bytes())?;
            file.write("\r\n".as_bytes())?;
        } 
        Ok(())
    }
}


fn compile(path: &str) -> std::io::Result<()> {
    // parser
    let parser = Parser::new(&path);
    let codes = parser.parse();
    // prepare file paths
    let outpath = path.replace("vm", "asm");
    let path_obj = Path::new(path);
    let fname = path_obj.file_stem().unwrap();
    // code generate
    CodeWriter::generate_code(fname.to_str().unwrap(), &outpath, codes)?;
    Ok(())
}

fn main() -> std::io::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 2 {
        println!("Usage: ./mvm /path/to/some.asm");
    } else {
        compile(&args[1])?;
    }
    Ok(())
}
